# -*- coding: utf-8 -*-

import tensorflow as tf
from vgg16 import vgg16_network
from scipy.special import expit
from scipy import stats
import cv2
import numpy as np
import os
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_softmax


def read_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image.astype(np.float32)
    image = (image - [104, 110, 109]) / 255
    assert isinstance(image, np.ndarray)
    image = image.reshape([1] + list(image.shape))
    return image


def show_predict(img_path, predict):
    # predict = expit(predict)
    # predict = stats.threshold(predict, 0.5, None, 0)
    # predict = cv2.resize(predict, (320, 180))

    # ppm_array = np.zeros(shape=(180, 320, 3), dtype=np.int8)
    # for y in range(0, 180):
    #     for x in range(0, 320):
    #         if predict[y, x] <= 0.5:
    #             ppm_array[y, x, :] = [64, 0, 0]
    #         elif 0.5 < predict[y, x] < 0.6:
    #             ppm_array[y, x, :] = [0, 0, 0]
    #         else:
    #             ppm_array[y, x, :] = [0, 0, 64]
    # cv2.imwrite("1.ppm", ppm_array)

    predict = np.reshape(predict, newshape=(180, 320, -1))
    min_value, max_value = np.min(predict), np.max(predict)
    print(min_value, max_value)
    img = ((predict - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    img = cv2.medianBlur(img, ksize=3)

    # src_img = cv2.imread(img_path)
    # show_image = np.zeros(shape=(180, 320, 3), dtype=np.uint8)
    # show_image = show_image + [0, 255, 255]
    # show_image = show_image.astype(dtype=np.uint8, copy=False)
    # tmp = np.zeros_like(show_image)
    # np.copyto(tmp, show_image, where=img > 0)
    #
    # cv2.addWeighted(tmp, 0.2, src_img, 0.8, 0, src_img)
    #
    # src_img = cv2.resize(src_img, (640, 360))
    cv2.imshow("src", cv2.imread(img_path))
    cv2.imshow("show", img)
    cv2.waitKey()
    # return src_img


def dense_crf_inference(image_path, predict):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    predict = expit(predict)
    predict = stats.threshold(predict, 0.5, None, 0)
    predict = cv2.resize(predict, (320, 180))
    tmp = np.zeros(shape=(2, 180, 320), dtype=predict.dtype)
    tmp[0, :, :] = predict
    tmp[1, :, :] = 1 - predict
    d = dcrf.DenseCRF2D(320, 180, 2)  # width, height, label
    u = unary_from_softmax(tmp)
    print(u.shape)
    d.setUnaryEnergy(u)
    d.addPairwiseBilateral(sxy=(20, 20), srgb=(15, 15, 15), rgbim=image, compat=10, kernel=dcrf.DIAG_KERNEL,
                           normalization=dcrf.NORMALIZE_SYMMETRIC)
    q = d.inference(5)
    test = np.array(q).reshape((2, 180, 320))[0]
    show_predict(image_path, test)
    # m = np.argmax(q, axis=0).reshape((180, 320))
    # print(m.shape)
    # print(np.sum(m))


def main():
    folder = "/home/lijun/Dataset/vehicle_segment/test"
    images_path = [os.path.join(folder, name) for name in os.listdir(folder)
                   if "jpg" in name and "_mask.jpg" not in name]
    image_path = "/home/lijun/Dataset/vehicle_segment/test/0000093.jpg"
    with tf.Graph().as_default():
        test_x_input = tf.placeholder(dtype=tf.float32, shape=[None, 180, 320, 3],
                                      name="test_x_input")
        test_x_output = vgg16_network(test_x_input)
        saver = tf.train.Saver()
        # cv2.namedWindow("show", cv2.WINDOW_NORMAL)
        with tf.Session() as sess:
            saver.restore(sess, "./model/vehicle_L2Loss.npk-2700")
            for image_path in images_path:
                img = read_image(image_path)
                predict = sess.run(test_x_output, feed_dict={test_x_input: img})
                predict = predict[0]
                dense_crf_inference(image_path, predict)
            # show_predict(image_path, predict)
            # save_path = "./test_results/{:>07d}.jpg".format(i)
            # show_img = cv2.resize(show_img, (640, 360))
            # cv2.imwrite(save_path, show_img)
            # cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
